"""
enrich完善问题示例，通过多轮对话，引导用户完善用户问题信息
"""
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate, SystemMessagePromptTemplate, MessagesPlaceholder, \
    HumanMessagePromptTemplate
from langchain_core.runnables import RunnableWithMessageHistory, RunnableConfig
from langchain_community.chat_message_histories import ChatMessageHistory
from models import get_ds_model_client
import langchain

langchain.debug=True

# 业务场景模板
templates = {
    "订机票": ["起点", "终点", "时间", "座位等级", "座位偏好"],
    "订酒店": ["城市", "入住时间", "离店时间", "房间数", "房型"]
}
prompt = PromptTemplate.from_template(
    "根据用户输入:```{user_input}```, 选择合适的业务模板。可用模板如下:{templates}, 直接返回业务模板的名词, 如'订机票'。")
llm = get_ds_model_client()
chain = prompt | llm

user_input = "我想订一张重庆到北京的机票"
intent = chain.invoke({"user_input": user_input, "templates": templates}).content
print("意图：", intent)

selected_template = templates.get(intent)
print("模板：", selected_template)

# 补充信息提示模板
info_template = f"""
    请根据用户原始问题和模板，判断原始问题是否完善。如果问题缺乏需要的信息，请生成一个友好的请求，明确指出需要补充的信息。若问题完善后，返回包含所有信息的完整问题。

### 原始问题
{user_input}

### 模板
{",".join(selected_template)}


### 输出示例
{{
    "isComplete": true,
    "content": "`完整问题`"
}}
{{
    "isComplete": false,
    "content": "`友好的引导到需要补充信息`"
}}
"""

chat_template = ChatPromptTemplate.from_messages(
    [SystemMessagePromptTemplate.from_template("你是一个信息补充助手，任务是分析用户问题需要的信息是否已完善，并返回json格式数据。"),
     ("placeholder", "{history}"),
     HumanMessagePromptTemplate.from_template("{input}")])
chat_chain = chat_template | llm
chat_message_history = ChatMessageHistory()

# 自动处理历史记录, 每次更新记录，但是暂时没有考虑到多用户场景
with_message_history = RunnableWithMessageHistory(chat_chain, lambda session_id: chat_message_history,
                                     input_messages_key="input", history_messages_key="history")

info_request = with_message_history.invoke({"input": info_template},
                                      config=RunnableConfig(configurable={"session_id": "unused"})).content
print("infor_request:", info_request)

parse = JsonOutputParser()
json_data = parse.parse(info_request)
print("json_data:", json_data)

while json_data.get("isComplete") is False:
    try:
        user_answer = input(f"\033[1;33m{json_data['content']}\033[0m\n你的回复：")
        info_request = with_message_history.invoke({"input": user_answer},
                                                   config=RunnableConfig(configurable={"session_id": "unused"})).content
        json_data = parse.parse(info_request)
        print("json_data:", json_data)
    except Exception as e:
        print(e)

# 输出最终结果
print(f"\033[1;32m[最终查询] {info_request}\033[0m")